# ==============================================================================
# Copyright (c) 2025 CompanyName. All rights reserved.
# Author:         22020873 陈泽欣
# Project:        Design of Deep Learning Fundamental Course
# Module:         identify_surface.py
# Date:           2025-05-24
# Description:    本模块用于识别骰子表面区域并提取关键角点信息，是视觉识别流程中的图像分割与特征提取核心组件。
#                 主要功能包括：
#                 - 基于 HSV 颜色空间的多色种分割（color_split）；
#                 - 图像颜色分析与命名（hsv_to_color_name）；
#                 - 提取轮廓角点以进行后续 PnP 计算（extract_corner_points）；
#                 - 根据颜色掩码查找并筛选有效物体轮廓（find_object_contours）；
#                 是整个骰子姿态识别系统中图像预处理与表面特征提取的关键模块。
# ==============================================================================

import cv2
import operator
from configs.norm_data import *

count = 0
def color_split(img, color):
    # 将图像从BGR转换为HSV色彩空间
    global count
    count += 1
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

    # 根据颜色定义HSV范围
    if color == 'red':
        mask1 = cv2.inRange(hsv, lower_red1, upper_red1)
        mask2 = cv2.inRange(hsv, lower_red2, upper_red2)
        mask = mask1 | mask2
    elif color == 'blue':
        mask = cv2.inRange(hsv, lower_blue, upper_blue)
    elif color == 'green':
        mask = cv2.inRange(hsv, lower_green, upper_green)
    elif color == "yellow":
        mask = cv2.inRange(hsv, lower_yellow, upper_yellow)
    elif color == "orange":
        mask = cv2.inRange(hsv, lower_orange, upper_orange)
    elif color == "cyan":
        mask = cv2.inRange(hsv, lower_cyan, upper_cyan)
    else:
        mask = np.zeros_like(hsv[:, :, 0])

    # kernel = np.ones((5, 5), np.uint8)  # 创建一个5x5的核
    # kernel2 = np.ones((5, 5), np.uint8)  # 创建一个5x5的核
    # mask = cv2.dilate(mask, kernel2, iterations=5)
    # mask = cv2.erode(mask, kernel, iterations=5)  # 腐蚀操作

    # if color == "green" or color == "cyan":s
    #     cv2.imwrite(f"green\\{count}.png", mask)
    # resized_img = cv2.resize(mask, None, fx=0.5, fy=0.5)
    # cv2.imshow(color, resized_img)

    return mask


def hsv_to_color_name(h, s, v):
    # 先检查明度（V）和饱和度（S），太低则归类为灰/黑/白
    if v < 0.2:
        return "Black"
    if s < 0.1:
        if v > 0.85:
            return "White"
        elif v > 0.5:
            return "Light Gray"
        else:
            return "Dark Gray"

    # 主要颜色判断（H: 0-360, S: 0-1, V: 0-1）
    if (h <= 10 or h >= 350):  # Red 范围（包含紫红过渡）
        if s > 0.7 and v > 0.7:
            return "red"
        else:
            return "red"
    elif h <= 25:  # Orange
        if s > 0.6 and v > 0.8:
            return "orange"
        else:
            return "orange"
    elif h <= 45:  # Yellow
        if s > 0.5 and v > 0.9:
            return "yellow"
        else:
            return "yellow"
    elif h <= 85:  # Green
        if s > 0.6 and v > 0.7:
            return "green"
        else:
            return "green"
    elif h <= 160:  # Cyan & Blue
        if h <= 100:  # Cyan
            return "cyan"
        else:  # Blue
            if s > 0.5 and v > 0.8:
                return "blue"
            else:
                return "blue"
    else:
        return "Undefined"  # 其他颜色（如紫色、粉色）不归类


def get_mode_channel(roi):
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)

    black_mask = cv2.inRange(hsv, lower_black, upper_black)

    # 得到非黑色区域的掩码
    non_black_mask = cv2.bitwise_not(black_mask)

    # 提取 H、S、V 三个通道
    h = hsv[:, :, 0]
    s = hsv[:, :, 1]
    v = hsv[:, :, 2]

    # 只统计非黑色区域的像素
    h_non_black = h[non_black_mask > 0]
    s_non_black = s[non_black_mask > 0]
    v_non_black = v[non_black_mask > 0]

    if len(h_non_black) == 0:
        return 0, 0, 0  # 如果没有非黑色像素，返回黑色

    # 分别统计各通道的众数
    def get_mode(arr):
        hist, _ = np.histogram(arr, bins=256, range=(0, 256))
        return np.argmax(hist)

    hue = int(get_mode(h_non_black))
    sat = int(get_mode(s_non_black))
    val = int(get_mode(v_non_black))

    return hue, sat, val

# 获取目标区域的角点
def extract_corner_points(contour):
    idx_br, _ = max(enumerate([pt[0][0] + pt[0][1] for pt in contour]), key=operator.itemgetter(1))
    idx_tl, _ = min(enumerate([pt[0][0] + pt[0][1] for pt in contour]), key=operator.itemgetter(1))
    idx_bl, _ = min(enumerate([pt[0][0] - pt[0][1] for pt in contour]), key=operator.itemgetter(1))
    idx_tr, _ = max(enumerate([pt[0][0] - pt[0][1] for pt in contour]), key=operator.itemgetter(1))
    return [contour[idx_tl][0], contour[idx_tr][0], contour[idx_br][0], contour[idx_bl][0]]
